import re
import json
import random
from langgraph.graph import END, StateGraph
from typing import Dict, TypedDict, List, Any, Literal, Optional, Callable

# 定义工具结构
class Tool:
    def __init__(self, name: str, description: str, func: Callable):
        self.name = name
        self.description = description
        self.func = func

    def __call__(self, **kwargs) -> Any:
        return self.func(**kwargs)

    def to_dict(self) -> Dict[str, Any]:
        return {
            "name": self.name,
            "description": self.description
        }

# 定义大模型响应结构
class ModelResponse(TypedDict):
    thought: str
    tool: Optional[str]  # 可能为None，表示不需要工具
    parameters: Dict[str, Any]

# 定义状态结构
class AgentState(TypedDict):
    user_input: str           # 用户输入的字符串
    model_response: ModelResponse  # 大模型结构化响应
    tool_result: Any          # 工具执行结果
    final_output: str         # 最终输出给用户的内容
    conversation: List[Dict]  # 对话历史

# 工具注册表
class ToolRegistry:
    def __init__(self):
        self.tools = {}

    def register(self, name: str, description: str):
        """工具注册装饰器"""
        def decorator(func):
            self.tools[name] = Tool(name, description, func)
            return func
        return decorator

    def get_tool(self, name: str) -> Optional[Tool]:
        return self.tools.get(name)

    def list_tools(self) -> List[Dict]:
        return [tool.to_dict() for tool in self.tools.values()]

# 创建工具注册表实例
tool_registry = ToolRegistry()

# 使用装饰器注册工具 - 就像定义普通函数一样
@tool_registry.register(name="add", description="计算两个数字的和")
def add(a: float, b: float) -> float:
    """加法函数 - 就像普通函数一样定义"""
    return a + b

@tool_registry.register(name="subtract", description="计算两个数字的差")
def subtract(a: float, b: float) -> float:
    """减法函数 - 就像普通函数一样定义"""
    return a - b

# 模拟大模型对话节点
def model_dialogue_node(state: AgentState) -> Dict:
    """模拟大模型对话节点，分析用户输入并决定使用哪个工具"""
    print(f"\n💬 对话历史: {state.get('conversation', [])[-1:] if state.get('conversation') else '无'}")
    print(f"🧠 用户输入: '{state['user_input']}'")
    print(f"🛠️ 可用工具: {[tool.name for tool in tool_registry.tools.values()]}")

    # 模拟大模型思考过程
    thought = ""
    tool_name = None
    parameters = {}

    # 提取数字
    numbers = re.findall(r'-?\d+\.?\d*', state["user_input"])  # 支持小数

    # 根据输入内容决定响应
    if "加" in state["user_input"] or "+" in state["user_input"] or "和" in state["user_input"]:
        thought = "用户要求进行加法运算"
        tool_name = "add"
    elif "减" in state["user_input"] or "-" in state["user_input"] or "差" in state["user_input"]:
        thought = "用户要求进行减法运算"
        tool_name = "subtract"
    elif any(char.isdigit() for char in state["user_input"]):
        # 包含数字但没有明确操作 - 需要澄清
        thought = "用户提供了数字但未明确操作类型，需要澄清"
        tool_name = None
    else:
        # 非数学问题
        thought = "这不是数学问题，将直接回复用户"
        tool_name = None

    # 处理需要工具的情况
    if tool_name and tool_name in tool_registry.tools:
        tool = tool_registry.get_tool(tool_name)

        # 尝试提取参数
        try:
            if len(numbers) >= 2:
                parameters = {"a": float(numbers[0]), "b": float(numbers[1])}
                thought += f"。将使用数字 {numbers[0]} 和 {numbers[1]}"
            else:
                thought += "。输入中缺少足够数字，将请求用户提供"
                tool_name = None
        except ValueError:
            thought += "。数字解析失败，将请求用户提供有效数字"
            tool_name = None

    # 构建结构化响应
    response: ModelResponse = {
        "thought": thought,
        "tool": tool_name,
        "parameters": parameters
    }

    # 添加随机性使模拟更真实
    if random.random() < 0.2:  # 20%概率添加额外思考
        response["thought"] += "（额外思考：用户可能需要更多帮助）"

    print(f"🤖 模拟大模型响应: {json.dumps(response, indent=2, ensure_ascii=False)}")

    # 更新对话历史
    conversation = state.get("conversation", [])
    conversation.append({
        "role": "assistant",
        "content": thought,
        "tool": tool_name
    })

    return {
        "model_response": response,
        "conversation": conversation
    }

# 通用工具执行函数
def execute_tool(state: AgentState) -> Dict:
    """执行模型选择的工具"""
    if not state["model_response"]["tool"]:
        return {"tool_result": None}  # 没有工具需要执行

    tool_name = state["model_response"]["tool"]
    tool = tool_registry.get_tool(tool_name)

    if not tool:
        print(f"⚠️ 错误: 工具 '{tool_name}' 未注册")
        return {"tool_result": f"错误: 工具 '{tool_name}' 未注册"}

    params = state["model_response"]["parameters"]

    print(f"🔧 执行工具: {tool_name} 参数: {params}")

    try:
        # 调用工具函数 - 就像调用普通函数一样
        result = tool(**params)
        print(f"⚙️ 工具结果: {result}")
        return {"tool_result": result}
    except Exception as e:
        error_msg = f"工具执行错误: {str(e)}"
        print(f"❌ {error_msg}")
        return {"tool_result": error_msg}

# 生成最终响应
def generate_final_response(state: AgentState) -> Dict:
    """生成最终输出给用户的内容"""
    tool_name = state["model_response"]["tool"]
    params = state["model_response"].get("parameters", {})
    result = state.get("tool_result")
    thought = state["model_response"]["thought"]

    # 根据情况生成不同响应
    if not tool_name:
        # 不需要工具的情况
        if "需要澄清" in thought:
            output = "您想进行加法还是减法运算？"
        elif "直接回复" in thought:
            output = "您好！我可以帮助您进行加法和减法运算。请告诉我您想计算什么？"
        else:
            output = "我无法理解您的请求，请提供更清晰的数学问题。"
    elif isinstance(result, str) and "错误" in result:
        output = f"计算失败: {result}"
    elif tool_name == "add":
        output = f"加法结果: {params.get('a', '?')} + {params.get('b', '?')} = {result}"
    elif tool_name == "subtract":
        output = f"减法结果: {params.get('a', '?')} - {params.get('b', '?')} = {result}"
    else:
        output = f"计算结果: {result}"

    # 更新对话历史
    conversation = state.get("conversation", [])
    conversation.append({
        "role": "assistant",
        "content": output
    })

    print(f"📝 最终输出: {output}")
    return {
        "final_output": output,
        "conversation": conversation
    }

# 构建图工作流
def create_workflow():
    # 初始化状态图
    workflow = StateGraph(AgentState)

    # 添加节点
    workflow.add_node("model_dialogue", model_dialogue_node)
    workflow.add_node("execute_tool", execute_tool)
    workflow.add_node("generate_response", generate_final_response)

    # 设置入口点
    workflow.set_entry_point("model_dialogue")

    # 添加条件边 - 根据模型决策路由
    def should_use_tool(state: AgentState) -> Literal["use_tool", "no_tool"]:
        """根据模型响应决定是否使用工具"""
        if state["model_response"]["tool"]:
            print("⤵️ 路由决策: 需要执行工具")
            return "use_tool"
        else:
            print("⤵️ 路由决策: 直接生成响应")
            return "no_tool"

    # 条件路由
    workflow.add_conditional_edges(
        "model_dialogue",
        should_use_tool,
        {
            "use_tool": "execute_tool",
            "no_tool": "generate_response"
        }
    )

    # 连接工具执行后的节点
    workflow.add_edge("execute_tool", "generate_response")

    # 连接到结束
    workflow.add_edge("generate_response", END)

    return workflow.compile()

# 执行示例
if __name__ == "__main__":
    # 初始化工作流
    app = create_workflow()

    # 打印注册的工具
    print("已注册工具:")
    for tool in tool_registry.list_tools():
        print(f"- {tool['name']}: {tool['description']}")

    # 测试案例
    test_cases = [
        {"input": "5加3等于多少？", "expected": ["5 + 3 = 8"]},
        {"input": "10减去7", "expected": ["10 - 7 = 3"]},
        {"input": "3和5", "expected": ["加法", "减法", "澄清"]},
        {"input": "你好", "expected": ["帮助", "加法", "减法"]},
        {"input": "计算圆的面积", "expected": ["无法理解", "数学问题"]},
    ]

    for test in test_cases:
        print("\n" + "="*50)
        print(f"测试输入: '{test['input']}'")

        # 初始状态
        state = {
            "user_input": test["input"],
            "conversation": [
                {"role": "user", "content": test["input"]}
            ]
        }

        result = app.invoke(state)

        # 验证结果
        final_output = result["final_output"]
        passed = any(exp in final_output for exp in test["expected"])

        if passed:
            print(f"✅ 测试通过! 输出: '{final_output}'")
        else:
            print(f"❌ 测试失败! 期望包含: {test['expected']}，实际: '{final_output}'")

        # 打印对话历史
        print("\n🗣️ 完整对话历史:")
        for msg in result["conversation"]:
            role = msg["role"].capitalize()
            content = msg["content"]
            tool = msg.get("tool", "")
            print(f"{role}: {content} {f'(工具: {tool})' if tool else ''}")
